# """
# A model for the C-P-U evolution 310-290 Ma
# Author: Shihan Li
# Several steps:
# 1. Auto spin to steady state in the beginning
# 2. Inverse to fit pco2 and d13c
# 3. For d13c, randomly assume the other end-member d13c for co2 source
# 4. Output the U isotope value, compare with proxy records

# Initial steady state:
# 1. t = 310Ma
# 2. pCO2 = 500e-6
# 3. o = 4.4e19 (after COPSE results)
# 4. d13c = 4.42 (after proxy records)
# 5. d235u = -0.10-0.27 (after proxy records)

# Forcing：
# 1. linear weatherability scale for silicate weathering
# 2. linear weatherability scale for carbonate weathering
# """

import numpy as np
from scipy.integrate import solve_ivp
from scipy import interpolate
from scipy.optimize import fmin_l_bfgs_b, fmin

import timeit
import pandas as pd
import matplotlib.pyplot as plt

import functions as fts               # functions for ODE
import emcee

from datetime import datetime

import os 
os.environ['OMP_NUM_THREADS'] = '1'

target = pd.read_csv('For_inv.csv')
'''--------------------------    Initial fluxes at t = 310ma -------------------------------- '''

# Constants
k_logistic = 12              # new determines slope of logistic anoxia function, COPSE_reloaded, Clarkson et al. 2018
k_uptake = 0.5000            # new determines efficiency of nutrient uptake in anoxia function; COPSE_reloaded; Clarkson et al. 2018
k_CtoK = 273.15    # convert degrees C to Kelvin
k_c = 4.3280                 # new determines climate sensitivity to CO2
k_l = 7.4000                 # new determines temperature sensitivity to luminosity
k_oxfrac = 0.9975            # updated initial oxic fraction
k_oceanmass = 1.397e21       # ocean mass (kg)

f_oxw_a = 0.5                # oxidative weathering dependency on O2 concentration
f_mocb_b = 2                 # marine organic carbon burial dependency on new production

a0 = 3.193e18       # atmosphere-ocean co2
o0 = 3.7e19         # atmosphere-ocean o2

pco2_i = 500e-6      
pco2atm0 = 280e-6
pco2pal_i = pco2_i/280e-6
a_i = a0 * np.sqrt(pco2pal_i)

delta_ocn_i = 4.46
delta_u_i = -0.10-0.27             # diagenetic correction, after Chen et al., 2022

k_oxidw = 5e12      # oxidative weatheirng + degassing, kept constant for simplicity
k_locb = 2.5e12     # continental organic C burial
k_mocb = 2.5e12     # marine organic C burial

# carbon isotope



delta_c = 4.5
delta_vol = -4
# organic weathering is kept free to close the 13c cycle

# organic carbon cycle follows the modern world
oxidw_i = k_oxidw
locb_i = k_locb
mocb_i = k_mocb

k6_fepb = 1e10               # updated Fe-P burial (mol/year)
k7_capb = 2e10               # updated Ca-P burial (mol/year)
k_mopb = 1e10                # organic-P burial in the ocean
k10_phosw = 4e10             # updated P weathering (mol/year)

newp0 = 117 * 2.2            # new production (umol/kg)
p0 = 2.2 * 1e-6 * k_oceanmass               # ocean (phosphate) phosphorus




# U cycle follows the modern value, after clarkson et al., 2018
u0 = 1.85e13         # modern U in the ocean

u_riv0 = 4.79e7       # river input
u_hydro0 = 5.7e6      # hydrothermal output
u_anox0 = 6.2e6       # anoxic sink
u_other0 = 3.6e7

u_i = u0
u_riv_i = u_riv0
u_hydro_i = u_hydro0
u_anox_i = u_anox0
u_other_i = u_other0

delta_u_riv = -0.29
d_u_hydro = 0.2
delta_u_hydro = delta_u_i + d_u_hydro
d_u_anox = 0.7
delta_u_anox = delta_u_i + d_u_anox

delta_u_other = (u_riv_i * delta_u_riv - delta_u_hydro * u_hydro_i - delta_u_anox*u_anox_i)/u_other_i

d_u_other = delta_u_other-delta_u_i


# isotope fractionations
# d238u_riv = -0.29
# u_frac_anox = 0.5
# u_frac_sed = 0.0156
# u_frac_hydro = 0.2

silw_i = 3*1e12              # 1.25*modern
carbw_i = 8*1e12            # 1*modern
        

###########       Parameters for inversion       ################


temp_i = 285
# 15+273.15         # K, 283-287

po2pal_i = 1         # 1.0 - 1.2

ppal_i = 1.2         # 1.0 - 1.5

scale_silw = [1.4, 2.5, 1.1]  # integrated long-term silicate weathering scale at t =  300, 290, 285 Ma
scale_locb = 1 # integrated long-term carbonate weathering modifier, change to organic burial later
scale_degassing = 0.7   # relative outgassing scale, 0.8-1.2
# scale_u_riv = 0.8   # relative d_u_riv modifier to fit the lont-term d238u trend, chagne to others later
scale_d13c_oxidw = 1  # relative oxidw modifier to fit the long-term d13c trend

alpha = 0.33             # co2 dependence
te = 34                 # e-folding temperature dependence of continental weathering, 5-50, after Krissansen-Totton et a., 2018

p_on_silw = 1          # P dependence on silw flux

cinput = np.array([5,4,6,8,20])   # *1e19 Carbon for 4 carbon emission events, corresponding time is manually defined
delta_cinput = np.array([-20,-6,-6, -6, -6])

params =[ ppal_i, scale_silw[0], scale_silw[1], scale_silw[2],
         scale_locb,
         scale_degassing, 
         # p_on_silw,
         cinput[0], cinput[1], cinput[2], cinput[3], cinput[4], 
         delta_cinput[0],delta_cinput[1],delta_cinput[2],delta_cinput[3],delta_cinput[4]]


raw_pco2 = pd.read_excel('raw_data.xlsx', sheet_name= 'pco2')
raw_u = pd.read_excel('raw_data.xlsx', sheet_name= 'u_raw')
raw_d13c = pd.read_excel('raw_data.xlsx', sheet_name= 'd13c_raw')

'''--------------------------    Probabilty function   -------------------------------- '''
def log_probability(params):
    lp = log_prior(params)
    
    if not np.isfinite(lp):
        return -np.inf
    
    else:
        global fscale_silw, fscale_locb,fscale_degassing,  ccdeg_i,o_i,p_i, fcinp, ANOX_i, phi_i, ppal_i, fscale_u_riv, fscale_d13c_oxidw, p_on_silw, fccinp
        
        cinput = np.zeros(5)
        delta_cinput = np.zeros(5)
        scale_silw = np.zeros(3)
        ppal_i, scale_silw[0], scale_silw[1], scale_silw[2],scale_locb, scale_degassing, cinput[0], cinput[1], cinput[2], cinput[3], cinput[4],delta_cinput[0],delta_cinput[1],delta_cinput[2],delta_cinput[3],delta_cinput[4]= params
        p_on_silw = 1
        scale_u_river = 1
        
        
        
       
    
        
       
        # interp_scale_locb.extend(scale_locb)
        
        
        
        fscale_silw = interpolate.interp1d([-310e6, -305e6, -295e6, -285e6],[1,scale_silw[0], scale_silw[1], scale_silw[2]], bounds_error = False, fill_value = 1)
        fscale_locb = interpolate.interp1d([-310e6,-285e6], [1,scale_locb], bounds_error = False, fill_value = 1)
        fscale_degassing = interpolate.interp1d([-310e6, -285e6],[1,scale_degassing], bounds_error = False, fill_value = 1)
        fscale_u_riv = interpolate.interp1d([-310e6, -285e6],[1,scale_u_riv], bounds_error = False, fill_value = 1)
        # fscale_d13c_oxidw = interpolate.interp1d([-310e6, -290e6],[1,scale_d13c_oxidw], bounds_error = False, fill_value = 1)
    
        cinput_age = [-304.79e6, -303.9e6, -302.30e6, -301.30e6, -298.32e6, -297.08e6, -295.13e6, -293.73e6, -292.05e6, -287.57e6]
        # cinput_age = [-304.3e6, -304.15e6, -299.19e6, -298.43e6, -297e6, -296.4e6, -295.73e6, -293.91e6]
        cinput_rate = np.array(cinput)* 1e19 /np.array([304.79e6 - 303.9e6, -301.30e6+302.30e6, -297.08e6+298.32e6, 295.13e6-293.73e6, 292.05e6-287.57e6 ])    
        # cinput_rate = np.array(cinput)* 1e18 /np.array([-304.15e6 + 304.3e6, -298.43e6+299.19e6, -296.4e6+297e6, 295.73e6-293.91e6])    
        
        fcinp = interpolate.interp1d(cinput_age, [cinput_rate[0], 0, cinput_rate[1], 0, cinput_rate[2],0, cinput_rate[3],0, cinput_rate[4],0], kind = 'zero', bounds_error = False, fill_value = 0)
        fccinp = interpolate.interp1d(cinput_age, [delta_cinput[0], 0, delta_cinput[1], 0, delta_cinput[2],0, delta_cinput[3],0, delta_cinput[4],0], kind = 'zero', bounds_error = False, fill_value = 0)
      
        t_eval = np.sort(target.age.values)
        
        
        ##############   initialize the model at t = 310Ma #################
        
        # Carbon
        global mccb_i
        mccb_i = silw_i + carbw_i    # total carbon burial to maintain the alkalinity balance, after COPSE
        o_i = o0 * po2pal_i
    
        
        ccdeg_i = silw_i
        
        # ap_i = ccdeg_i+oxidw_i-locb_i-mocb_i+carbw_i-mccb_i       # check the balance of carbon cycle
        
        # C isotope
        phi_i = 0.01614 * (a_i/a0)  # fraction of C in atmosphere:ocean
        
        d_locb_i, D_P_i, d_mocb_i, D_B_i, d_mccb_i, d_ocean_i, d_atmos_i = fts.Cisotopefrac(temp_i, pco2pal_i, po2pal_i, phi_i)
        
        
        delta_a_i = delta_ocn_i - d_ocean_i
        
        delta_locb = delta_a_i+d_locb_i
        delta_mocb = delta_a_i+d_mocb_i
        
        
        delta_mccb_i = delta_a_i + d_mccb_i
        
        global delta_g
        delta_g = (locb_i * ( delta_locb) + mocb_i * (delta_mocb) - ccdeg_i * delta_vol - carbw_i * delta_c + mccb_i * delta_mccb_i)/oxidw_i
        
        moldelta_a_i = a_i * delta_a_i
        
        
        # moldelta_ap_i = -ccdeg_i*delta_vol - oxidw_i*delta_g + locb_i*delta_locb + mocb_i*delta_mocb - carbw_i*delta_c + mccb_i*delta_mccb_i 
        
        
        # P cycle
        global p_i, newp_i, ANOW_i, mopb_i, fepb_i, capb_i, phosw_i
        p_i = p0 * ppal_i
        newp_i = 117 * (p_i/p0) * 2.2
        ANOX_i = 1/(1+np.exp(-k_logistic * (k_uptake * (newp_i/newp0)-po2pal_i)))   
        mopb_i = mocb_i * ((ANOX_i/1000)+((1-ANOX_i)/250))  # ocean burial
        fepb_i = (k6_fepb/k_oxfrac)*(1-ANOX_i)*(p_i/p0)
        capb_i = k7_capb * ((newp_i/newp0)**f_mocb_b)
        
        phosw_i = mopb_i + fepb_i + capb_i
        
        pp_i = phosw_i - mopb_i -fepb_i - capb_i
        
        # U cycle
        # up = u_riv_i - u_hydro_i - u_anox_i - u_other_i
        # moldelta_up_i = u_riv_i*d238u_riv - u_hydro_i*delta_u_hydro - u_anox_i*delta_u_anox - u_other_i*delta_u_other
        # print(up)
        # print(moldelta_up_i)
        moldelta_u_i = u_i * delta_u_i
        
        ystart = np.array([a_i, p_i, o_i, moldelta_a_i, u_i, moldelta_u_i])
        t0 = -310e6
        tfinal = -285e6
    
    
    
        start_time = timeit.default_timer()
        
        ysol = solve_ivp(derivs,(t0,tfinal), ystart, args={1}, t_eval = t_eval, method = 'LSODA',  max_step = 1e4)
        # ysol = derivs(-310e6, ystart, 1)
        
        # print("\n@ Starting integration")
        # print("[tstart tfinal]=[%.2e %.2e]\n" % (t0, tfinal))
        
        if np.isnan(ysol.y).any():
            return -np.inf
        
        else:
            t = ysol.t                # time
            y = ysol.y                # tracers

            nstep = len(t)

            params = np.zeros((nstep, 14))

            for i in range(nstep):
                params[i,:] = derivs(t[i], y[:,i], 0)

            df_params = pd.DataFrame(params)
            df_params.columns=['Temperature','ccdeg', 'oxidw', 'locb', 'mocb', 'silw', 'carbw', 'mccb', 'delta_ocn', 'phosw', 'mopb', 'fepb', 'capb', 'ANOX']
            df_params['Age'] = t

            df_sol = pd.DataFrame(ysol.y.T)
            df_sol.columns=['A',  'P',  'O', 'moldelta_A', 'U', 'moldelta_U']


            df_sol['Age'] = ysol.t
            df_sol['phosphate_m'] = (df_sol['P']/k_oceanmass) * 1e6  # umol/kg
            df_sol['p/p0'] = df_sol['phosphate_m']/2.2
            df_sol['U_m'] = (df_sol['U']/k_oceanmass)*1e6            # umol/kg
            df_sol['d235U'] = (df_sol['moldelta_U']/df_sol['U'])     # d235U
            df_sol['CO2_PAL'] = (df_sol['A']/a_i)**2
            df_sol['d13c'] = (df_sol['moldelta_A']/df_sol['A'])      # d13c

            df_sol['CO2atm'] = df_sol['CO2_PAL'] * pco2_i * 1e6
            df_sol['O2PAL'] = df_sol['O']/o0
            # df_sol.to_csv("tracer.csv")
            # df_params.to_csv("parameters.csv")

#             fig, axes = plt.subplots(figsize = (12,10), nrows = 3, ncols = 3)
            
#             raw_pco2.plot(x='age', y='pco2', ax=axes[0,0], marker = '*', lw=0)
#             df_sol.plot(x='Age', y='CO2atm', ax=axes[0,0])
            
#             df_sol.plot(x='Age', y='U_m', ax=axes[0,1])
            
#             raw_u.plot(x='age', y = 'u', ax = axes[0,2], marker='*', lw=1.5)
#             df_sol.plot(x='Age', y='d235U', ax=axes[0,2])
            
#             df_sol.plot(x='Age', y='phosphate_m', ax=axes[1,0])
#             df_params.plot(x='Age', y ='ANOX', ax=axes[1,1])
#             # axes[1,2].remove()
#             # df_params.plot(x='Age', y ='oxidw', ax=axes[1,2])
#             # df_params.plot(x='Age', y ='locb', ax=axes[2,0])
#             # df_params.plot(x='Age', y ='mocb', ax=axes[2,1])
#             df_params.plot(x='Age', y ='silw', ax=axes[1,2])
#             # df_params.plot(x='Age', y ='carbw', ax=axes[2,2])
            
#             raw_d13c.plot(x='age', y='d13c', ax=axes[2,0], marker = '*', lw=1.5)
#             df_sol.plot(x='Age', y='d13c', ax=axes[2,0])
            
#             df_sol.plot(x='Age', y='O2PAL', ax=axes[2,1])
            
#             axes[2,2].plot(t, fcinp(t)/1e15)
           
#             axes[2,2].set_ylabel('Carbon emission (Gt/yr)')


#             plt.tight_layout()

            pco2_proxy = target['pco2'].values
            pco2_std = target['pco2_std'].values
            d13c_proxy = target['d13c'].values
            d13c_std = target['d13c_std'].values
            u_proxy = target['u'].values
            u_std = target['u_std'].values

            pco2_model = df_sol['CO2atm'].values
            d13c_model = df_sol['d13c'].values
            u_model = df_sol['d235U'].values

            # u_index1 = range(21)
            # u_index2 = 35
            
            # index = [0,5,11,15,20,37, 44,50,55,60,65,70,73]

            sum_diff = sum((pco2_proxy-pco2_model)**2/(pco2_std**2)) +   sum((u_proxy[0:-1]-u_model[0:-1])**2/(u_std[0:-1]**2)) + sum((d13c_proxy[0:-1]-d13c_model[0:-1])**2/(d13c_proxy[0:-1]**2))
            sum_diff += sum(np.log(2*np.pi*pco2_std**2))+sum(np.log(2*np.pi*u_std[0:-1]**2))+sum(np.log(2*np.pi*d13c_std[0:-1]**2))
            
            print(-0.5*sum_diff)

            return lp-0.5 * sum_diff
            # return 0



def derivs(t, y, switch):
    # if t<-310e6:
    #     t = -310e6

    # t = -310e6
    t_Ma = t/1e6

  


    # retrieve the parameters
    a, p, o, moldelta_a, u, moldelta_u = y

    delta_a = moldelta_a/a
    delta_u = moldelta_u/u
    # calcualte pco2
    po2pal = o/o0

    pco2pal = (a/a_i)**2
    pco2 = pco2_i * pco2pal
    phi = 0.01614 * (a/a_i)

    # temp = temp_i
    temp = temp_i + k_c * np.log(pco2pal) + k_l/570e6 * (t+310e6)
    
    diff_temp = temp - temp_i

    """---------------------  Carbon  -------------------------"""
    
    # degassing
    ccdeg = ccdeg_i * fscale_degassing(t)
   
    silw = fscale_silw(t) * silw_i * (pco2pal ** alpha) * np.exp(diff_temp/te) 
    
    
    
    
    carbw = fscale_silw(t) * carbw_i * (pco2pal ** alpha) * np.exp(diff_temp/te)          

    # oxidw
    oxw_fac = po2pal ** f_oxw_a
    oxidw = oxidw_i *  oxw_fac 

    

    # burial
    mccb = silw + carbw
    locb =  fscale_locb(t) * locb_i * (2*pco2pal/(1+pco2pal))
    mocb = mocb_i * (p/p_i) ** f_mocb_b

    ap = ccdeg + oxidw  - locb - mocb - silw + fcinp(t)/12
    
    
    
   
   

    """--------------------- C isotope  -------------------------"""
    d_locb, D_P, d_mocb, D_B, d_mccb, d_ocean, d_atmos = fts.Cisotopefrac(temp, pco2pal * pco2pal_i, po2pal*po2pal_i, phi * np.sqrt(pco2pal_i))

    delta_mccb = delta_a + d_mccb
    delta_ocn = delta_a + d_ocean
    delta_mocb = delta_a + d_mocb
    delta_locb = delta_a + d_locb
    # moldelta_ap =  -locb*(delta_a+d_locb) - mocb * (delta_a+d_mocb) + oxidw*delta_g  + ccdeg*delta_vol + carbw*delta_c - mccb*delta_mccb
    moldelta_ap =  -locb*(delta_locb) - mocb * (delta_mocb) + oxidw*delta_g + ccdeg*delta_vol + carbw*delta_c - mccb*delta_mccb + fcinp(t)/12 * fccinp(t)
    
    

    """--------------------- P  -------------------------"""
    # P cycle
    Pconc = (p/p0) * 2.2

    # marine new production
    newp = 117 * Pconc
    # anoxic
    ANOX =  1/(1+np.exp(-k_logistic * (k_uptake * (newp/newp0)-po2pal)))
    
    # phosw_i = k_phosw * silw_i/k_silw
    mopb = mocb * ((ANOX/1000)+((1-ANOX)/250))

    fepb = (k6_fepb/k_oxfrac)*(1-ANOX)*(p/p0)
    # fepb_i = (k_fepb/k_oxfrac)*(1-ANOX_i)
    capb = k7_capb * ((newp/newp0)**f_mocb_b)
    # capb_i = k_capb * ((newp_i/newp0))

    # phosphorous balance
    phosw = phosw_i * ((silw)/(silw_i)) ** p_on_silw


    pp = phosw-mopb-fepb-capb
    
    """--------------------- O  -------------------------"""
    op = locb + mocb - oxidw 
    
    # U cycle
    u_riv = u_riv_i * (silw/silw_i) 
    u_hydro = fscale_degassing(t) *u_hydro_i 
    u_anox = u_anox_i * (ANOX/ANOX_i) * u/u_i
    
    
    
    u_other =  u_other_i * (u/u_i)
    


    moldelta_up = u_riv * delta_u_riv * fscale_u_riv(t) - u_hydro*(delta_u+d_u_hydro) - u_anox*(delta_u+d_u_anox)- u_other*(delta_u+d_u_other)

    up = u_riv - u_hydro - u_anox - u_other



    
    if switch:
        yp = np.array([ap, pp, op, moldelta_ap, up, moldelta_up])
        return yp
    else:
        params = np.array([temp, ccdeg, oxidw, locb, mocb, silw, carbw, mccb, delta_ocn, phosw, mopb, fepb, capb, ANOX])
        return params

def log_prior(theta):
    
    cinput = np.zeros(5)
    delta_cinput = np.zeros(5)
    scale_silw = np.zeros(3)
    ppal_i, scale_silw[0], scale_silw[1], scale_silw[2], scale_locb, scale_degassing,  cinput[0], cinput[1], cinput[2], cinput[3], cinput[4],delta_cinput[0],delta_cinput[1],delta_cinput[2],delta_cinput[3],delta_cinput[4]= theta
    if 1.0<=ppal_i<=1.5 and np.all(scale_silw>=0.8) and np.all(scale_silw<=3) and 0.9<= scale_locb<=1.1 and 0.7<=scale_degassing<=1.3 and 0.5<=cinput[0]<=10 and 0.5<=cinput[1]<=10 and 0.5<=cinput[2]<=10 and .5<=cinput[3]<=10 and .5<=cinput[4]<=40 and np.all(delta_cinput>=-25) and np.all(delta_cinput<=0)  :
        return 0.0
    # return 0
    return -np.inf

from multiprocessing import Pool
params_test = [  1.23984592,   1.897619,   2.5,   1.059425,   1.07638667,
   0.86314637,      
#                0,0,0,0,0,
    6,   5.16283163,   6.44446281, 7.85036545,  30,
    -19.71662526, -2,  -6.75820334, -6.00757177,  -6.02946557]
params_init = params_test
with Pool() as pool:
                                                                                    
    ndim = len(params_init)

    pos = np.array(params_init) +  1e-1 * np.random.randn(50, ndim)
    nwalkers, ndim = pos.shape

    sampler = emcee.EnsembleSampler(
        nwalkers, ndim, log_probability
    )
    sampler.run_mcmc(pos, 5000, progress=True)

-59.88017721946259
-62.80280965380545
-64.94616630601458
-63.37595263148728
-63.69763683621285
-61.9493111044654
-60.39507912325345
-63.89479829156497
-61.4881480313158
emcee: Exception while calling your likelihood function:
  params: [  1.29547923   1.7943043    2.60512591   1.16098585   1.07501316
   0.87304223   5.98787476   5.11484296   6.33221024   7.84727101
  30.11402351 -19.56304348  -2.04849408  -6.89476631  -6.06448985
  -6.22686468]
  args: []
  kwargs: {}
  exception:

Traceback (most recent call last):
  File "C:\ProgramData\Anaconda3\lib\site-packages\emcee\ensemble.py", line 624, in __call__
    return self.f(x, *self.args, **self.kwargs)
  File "C:\Users\shihan\AppData\Local\Temp/ipykernel_9572/604939613.py", line 233, in log_probability
    ysol = solve_ivp(derivs,(t0,tfinal), ystart, args={1}, t_eval = t_eval, method = 'LSODA',  max_step = 1e4)
  File "C:\ProgramData\Anaconda3\lib\site-packages\scipy\integrate\_ivp\ivp.py", line 634, in solve_ivp
    ys.append(sol(t_eval_step))
  File "C:\ProgramData\Anaconda3\lib\site-packages\scipy\integrate\_ivp\base.py", line 252, in __call__
    return self._call_impl(t)
  File "C:\ProgramData\Anaconda3\lib\site-packages\scipy\integrate\_ivp\lsoda.py", line 188, in _call_impl
    return np.dot(self.yh, x)
  File "<__array_function__ internals>", line 5, in dot
KeyboardInterrupt

---------------------------------------------------------------------------
KeyboardInterrupt                         Traceback (most recent call last)
~\AppData\Local\Temp/ipykernel_9572/203772687.py in <module>
     16         nwalkers, ndim, log_probability
     17     )
---> 18     sampler.run_mcmc(pos, 5000, progress=True)
     19 

C:\ProgramData\Anaconda3\lib\site-packages\emcee\ensemble.py in run_mcmc(self, initial_state, nsteps, **kwargs)
    441 
    442         results = None
--> 443         for results in self.sample(initial_state, iterations=nsteps, **kwargs):
    444             pass
    445 

C:\ProgramData\Anaconda3\lib\site-packages\emcee\ensemble.py in sample(self, initial_state, log_prob0, rstate0, blobs0, iterations, tune, skip_initial_state_check, thin_by, thin, store, progress, progress_kwargs)
    342             state.blobs = blobs0
    343         if state.log_prob is None:
--> 344             state.log_prob, state.blobs = self.compute_log_prob(state.coords)
    345         if np.shape(state.log_prob) != (self.nwalkers,):
    346             raise ValueError("incompatible input dimensions")

C:\ProgramData\Anaconda3\lib\site-packages\emcee\ensemble.py in compute_log_prob(self, coords)
    487             else:
    488                 map_func = map
--> 489             results = list(map_func(self.log_prob_fn, p))
    490 
    491         try:

C:\ProgramData\Anaconda3\lib\site-packages\emcee\ensemble.py in __call__(self, x)
    622     def __call__(self, x):
    623         try:
--> 624             return self.f(x, *self.args, **self.kwargs)
    625         except:  # pragma: no cover
    626             import traceback

~\AppData\Local\Temp/ipykernel_9572/604939613.py in log_probability(params)
    231         start_time = timeit.default_timer()
    232 
--> 233         ysol = solve_ivp(derivs,(t0,tfinal), ystart, args={1}, t_eval = t_eval, method = 'LSODA',  max_step = 1e4)
    234         # ysol = derivs(-310e6, ystart, 1)
    235 

C:\ProgramData\Anaconda3\lib\site-packages\scipy\integrate\_ivp\ivp.py in solve_ivp(fun, t_span, y0, method, t_eval, dense_output, events, vectorized, args, **options)
    632                     sol = solver.dense_output()
    633                 ts.append(t_eval_step)
--> 634                 ys.append(sol(t_eval_step))
    635                 t_eval_i = t_eval_i_new
    636 

C:\ProgramData\Anaconda3\lib\site-packages\scipy\integrate\_ivp\base.py in __call__(self, t)
    250         if t.ndim > 1:
    251             raise ValueError("`t` must be a float or a 1-D array.")
--> 252         return self._call_impl(t)
    253 
    254     def _call_impl(self, t):

C:\ProgramData\Anaconda3\lib\site-packages\scipy\integrate\_ivp\lsoda.py in _call_impl(self, t)
    186             x = ((t - self.t) / self.h) ** self.p[:, None]
    187 
--> 188         return np.dot(self.yh, x)

<__array_function__ internals> in dot(*args, **kwargs)

KeyboardInterrupt: 

ndim = len(params_init)
fig, axes = plt.subplots(ndim, figsize=(12, 50), sharex=True)
samples = sampler.get_chain()
labels = ["ppal_i", "scale_silw1", "scale_silw2","scale_silw3", "scale_locb", 'scale_degassing', 'scale_u_riv', 'cinput1', 'cinput2', 'cinput3', 'cinput4', 'cinput5', 'd13c1', 'd13c2', 'd13c3', 'd13c4', 'd13c5']
for i in range(ndim):
    ax = axes[i]
    ax.plot(samples[:, :, i], "k", alpha=0.3)
    ax.set_xlim(0, len(samples))
    ax.set_ylabel(labels[i])
    ax.yaxis.set_label_coords(-0.1, 0.5)

axes[-1].set_xlabel("step number");

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
~\AppData\Local\Temp/ipykernel_9572/3182270289.py in <module>
      1 ndim = len(params_init)
      2 fig, axes = plt.subplots(ndim, figsize=(12, 50), sharex=True)
----> 3 samples = sampler.get_chain()
      4 labels = ["ppal_i", "scale_silw1", "scale_silw2","scale_silw3", "scale_locb", 'scale_degassing', 'scale_u_riv', 'cinput1', 'cinput2', 'cinput3', 'cinput4', 'cinput5', 'd13c1', 'd13c2', 'd13c3', 'd13c4', 'd13c5']
      5 for i in range(ndim):

C:\ProgramData\Anaconda3\lib\site-packages\emcee\ensemble.py in get_chain(self, **kwargs)
    580 
    581     def get_chain(self, **kwargs):
--> 582         return self.get_value("chain", **kwargs)
    583 
    584     get_chain.__doc__ = Backend.get_chain.__doc__

C:\ProgramData\Anaconda3\lib\site-packages\emcee\ensemble.py in get_value(self, name, **kwargs)
    600 
    601     def get_value(self, name, **kwargs):
--> 602         return self.backend.get_value(name, **kwargs)
    603 
    604     def get_autocorr_time(self, **kwargs):

C:\ProgramData\Anaconda3\lib\site-packages\emcee\backends\backend.py in get_value(self, name, flat, thin, discard)
     42     def get_value(self, name, flat=False, thin=1, discard=0):
     43         if self.iteration <= 0:
---> 44             raise AttributeError(
     45                 "you must run the sampler with "
     46                 "'store == True' before accessing the "

AttributeError: you must run the sampler with 'store == True' before accessing the results

params

tau = sampler.get_autocorr_time()
print(tau)

---------------------------------------------------------------------------
AutocorrError                             Traceback (most recent call last)
~\AppData\Local\Temp/ipykernel_9572/826269742.py in <module>
----> 1 tau = sampler.get_autocorr_time()
      2 print(tau)

C:\ProgramData\Anaconda3\lib\site-packages\emcee\ensemble.py in get_autocorr_time(self, **kwargs)
    603 
    604     def get_autocorr_time(self, **kwargs):
--> 605         return self.backend.get_autocorr_time(**kwargs)
    606 
    607     get_autocorr_time.__doc__ = Backend.get_autocorr_time.__doc__

C:\ProgramData\Anaconda3\lib\site-packages\emcee\backends\backend.py in get_autocorr_time(self, discard, thin, **kwargs)
    148         """
    149         x = self.get_chain(discard=discard, thin=thin)
--> 150         return thin * autocorr.integrated_time(x, **kwargs)
    151 
    152     @property

C:\ProgramData\Anaconda3\lib\site-packages\emcee\autocorr.py in integrated_time(x, c, tol, quiet)
    110         msg += "N/{0} = {1:.0f};\ntau: {2}".format(tol, n_t / tol, tau_est)
    111         if not quiet:
--> 112             raise AutocorrError(tau_est, msg)
    113         logger.warning(msg)
    114 

AutocorrError: The chain is shorter than 50 times the integrated autocorrelation time for 16 parameter(s). Use this estimate with caution and run a longer chain!
N/50 = 100;
tau: [465.5876854  419.54221881 428.77571678 366.65530337 438.99613223
 410.96065388 421.50406397 421.37374338 430.35296892 438.64270155
 408.46813186 430.64292991 537.51263278 396.10657625 443.04648856
 524.35672619]

flat_samples = sampler.get_chain(discard=100, thin=16, flat=True)
print(flat_samples.shape)

(15300, 16)

import corner

fig = corner.corner(
    flat_samples)

from IPython.display import display, Math
param_50_percent = np.zeros(ndim)
for i in range(ndim):
    mcmc = np.percentile(flat_samples[:, i], [16, 50, 84])
    param_50_percent[i] = mcmc[1]
    q = np.diff(mcmc)
    txt = "\mathrm{{{3}}} = {0:.3f}_{{-{1:.3f}}}^{{{2:.3f}}}"
    txt = txt.format(mcmc[1], q[0], q[1], labels[i])
    display(Math(txt))

$\displaystyle \mathrm{ppal_i} = 1.248_{-0.176}^{0.188}$

$\displaystyle \mathrm{scale_silw1} = 1.443_{-0.244}^{0.347}$

$\displaystyle \mathrm{scale_silw2} = 1.998_{-0.366}^{0.410}$

$\displaystyle \mathrm{scale_silw3} = 1.364_{-0.331}^{0.570}$

$\displaystyle \mathrm{scale_locb} = 0.987_{-0.060}^{0.079}$

$\displaystyle \mathrm{scale_degassing} = 1.013_{-0.200}^{0.200}$

$\displaystyle \mathrm{scale_u_riv} = 4.813_{-2.358}^{2.440}$

$\displaystyle \mathrm{cinput1} = 5.482_{-2.366}^{2.616}$

$\displaystyle \mathrm{cinput2} = 5.637_{-3.414}^{2.681}$

$\displaystyle \mathrm{cinput3} = 5.913_{-3.301}^{2.620}$

$\displaystyle \mathrm{cinput4} = 30.577_{-8.610}^{6.389}$

$\displaystyle \mathrm{cinput5} = -17.803_{-4.574}^{7.899}$

$\displaystyle \mathrm{d13c1} = -5.955_{-10.707}^{4.246}$

$\displaystyle \mathrm{d13c2} = -7.027_{-7.343}^{4.496}$

$\displaystyle \mathrm{d13c3} = -8.555_{-9.044}^{5.677}$

$\displaystyle \mathrm{d13c4} = -6.690_{-7.668}^{4.087}$

raw_pco2 = pd.read_excel('raw_data.xlsx', sheet_name= 'pco2')
raw_u = pd.read_excel('raw_data.xlsx', sheet_name= 'u_raw')
raw_d13c = pd.read_excel('raw_data.xlsx', sheet_name= 'd13c_raw')
po2pal_i = 1.0
def ode_plot(params):
    global fscale_silw, fscale_locb,fscale_degassing,  ccdeg_i,o_i,p_i, fcinp, ANOX_i, phi_i, ppal_i, fscale_u_riv, fscale_d13c_oxidw, p_on_silw, fccinp
        
    cinput = np.zeros(5)
    delta_cinput = np.zeros(5)
    scale_silw = np.zeros(3)
    ppal_i, scale_silw[0], scale_silw[1], scale_silw[2],scale_locb, scale_degassing, cinput[0], cinput[1], cinput[2], cinput[3], cinput[4],delta_cinput[0],delta_cinput[1],delta_cinput[2],delta_cinput[3],delta_cinput[4]= params
    p_on_silw = 1.0

    scale_u_river = 1






    # interp_scale_locb.extend(scale_locb)



    fscale_silw = interpolate.interp1d([-310e6, -305e6, -295e6, -285e6],[1,scale_silw[0], scale_silw[1], scale_silw[2]], bounds_error = False, fill_value = 1)
    fscale_locb = interpolate.interp1d([-310e6,-285e6], [1,scale_locb], bounds_error = False, fill_value = 1)
    fscale_degassing = interpolate.interp1d([-310e6, -285e6],[1,scale_degassing], bounds_error = False, fill_value = 1)
    fscale_u_riv = interpolate.interp1d([-310e6, -285e6],[1,scale_u_riv], bounds_error = False, fill_value = 1)
    # fscale_d13c_oxidw = interpolate.interp1d([-310e6, -290e6],[1,scale_d13c_oxidw], bounds_error = False, fill_value = 1)

    cinput_age = [-304.79e6, -303.9e6, -302.30e6, -301.30e6, -298.32e6, -297.08e6, -295.13e6, -293.73e6, -292.05e6, -287.57e6]
    # cinput_age = [-304.3e6, -304.15e6, -299.19e6, -298.43e6, -297e6, -296.4e6, -295.73e6, -293.91e6]
    cinput_rate = np.array(cinput)* 1e19 /np.array([304.79e6 - 303.9e6, -301.30e6+302.30e6, -297.08e6+298.32e6, 295.13e6-293.73e6, 292.05e6-287.57e6 ])    
    # cinput_rate = np.array(cinput)* 1e18 /np.array([-304.15e6 + 304.3e6, -298.43e6+299.19e6, -296.4e6+297e6, 295.73e6-293.91e6])    

    fcinp = interpolate.interp1d(cinput_age, [cinput_rate[0], 0, cinput_rate[1], 0, cinput_rate[2],0, cinput_rate[3],0, cinput_rate[4],0], kind = 'zero', bounds_error = False, fill_value = 0)
    fccinp = interpolate.interp1d(cinput_age, [delta_cinput[0], 0, delta_cinput[1], 0, delta_cinput[2],0, delta_cinput[3],0, delta_cinput[4],0], kind = 'zero', bounds_error = False, fill_value = 0)

    t_eval = np.sort(target.age.values)


    ##############   initialize the model at t = 310Ma #################

    # Carbon
    global mccb_i
    mccb_i = silw_i + carbw_i    # total carbon burial to maintain the alkalinity balance, after COPSE
    o_i = o0 * po2pal_i


    ccdeg_i = silw_i

    # ap_i = ccdeg_i+oxidw_i-locb_i-mocb_i+carbw_i-mccb_i       # check the balance of carbon cycle

    # C isotope
    phi_i = 0.01614 * (a_i/a0)  # fraction of C in atmosphere:ocean

    d_locb_i, D_P_i, d_mocb_i, D_B_i, d_mccb_i, d_ocean_i, d_atmos_i = fts.Cisotopefrac(temp_i, pco2pal_i, po2pal_i, phi_i)


    delta_a_i = delta_ocn_i - d_ocean_i

    delta_locb = delta_a_i+d_locb_i
    delta_mocb = delta_a_i+d_mocb_i


    delta_mccb_i = delta_a_i + d_mccb_i

    global delta_g
    delta_g = (locb_i * ( delta_locb) + mocb_i * (delta_mocb) - ccdeg_i * delta_vol - carbw_i * delta_c + mccb_i * delta_mccb_i)/oxidw_i

    moldelta_a_i = a_i * delta_a_i


    # moldelta_ap_i = -ccdeg_i*delta_vol - oxidw_i*delta_g + locb_i*delta_locb + mocb_i*delta_mocb - carbw_i*delta_c + mccb_i*delta_mccb_i 


    # P cycle
    global p_i, newp_i, ANOW_i, mopb_i, fepb_i, capb_i, phosw_i
    p_i = p0 * ppal_i
    newp_i = 117 * (p_i/p0) * 2.2
    ANOX_i = 1/(1+np.exp(-k_logistic * (k_uptake * (newp_i/newp0)-po2pal_i)))   
    mopb_i = mocb_i * ((ANOX_i/1000)+((1-ANOX_i)/250))  # ocean burial
    fepb_i = (k6_fepb/k_oxfrac)*(1-ANOX_i)*(p_i/p0)
    capb_i = k7_capb * ((newp_i/newp0)**f_mocb_b)

    phosw_i = mopb_i + fepb_i + capb_i

    pp_i = phosw_i - mopb_i -fepb_i - capb_i

    # U cycle
    # up = u_riv_i - u_hydro_i - u_anox_i - u_other_i
    # moldelta_up_i = u_riv_i*d238u_riv - u_hydro_i*delta_u_hydro - u_anox_i*delta_u_anox - u_other_i*delta_u_other
    # print(up)
    # print(moldelta_up_i)
    moldelta_u_i = u_i * delta_u_i

    ystart = np.array([a_i, p_i, o_i, moldelta_a_i, u_i, moldelta_u_i])
    t0 = -310e6
    tfinal = -285e6



    start_time = timeit.default_timer()

    ysol = solve_ivp(derivs,(t0,tfinal), ystart, args={1}, method = 'LSODA',  max_step = 1e4)
    # ysol = derivs(-310e6, ystart, 1)

    # print("\n@ Starting integration")
    # print("[tstart tfinal]=[%.2e %.2e]\n" % (t0, tfinal))

    if np.isnan(ysol.y).any():
        return -np.inf

    else:
        t = ysol.t                # time
        y = ysol.y                # tracers

        nstep = len(t)

        params = np.zeros((nstep, 14))

        for i in range(nstep):
            params[i,:] = derivs(t[i], y[:,i], 0)

        df_params = pd.DataFrame(params)
        df_params.columns=['Temperature','ccdeg', 'oxidw', 'locb', 'mocb', 'silw', 'carbw', 'mccb', 'delta_ocn', 'phosw', 'mopb', 'fepb', 'capb', 'ANOX']
        df_params['Age'] = t

        df_sol = pd.DataFrame(ysol.y.T)
        df_sol.columns=['A',  'P',  'O', 'moldelta_A', 'U', 'moldelta_U']


        df_sol['Age'] = ysol.t
        df_sol['phosphate_m'] = (df_sol['P']/k_oceanmass) * 1e6  # umol/kg
        df_sol['p/p0'] = df_sol['phosphate_m']/2.2
        df_sol['U_m'] = (df_sol['U']/k_oceanmass)*1e6            # umol/kg
        df_sol['d235U'] = (df_sol['moldelta_U']/df_sol['U'])     # d235U
        df_sol['CO2_PAL'] = (df_sol['A']/a_i)**2
        df_sol['d13c'] = (df_sol['moldelta_A']/df_sol['A'])      # d13c

        df_sol['CO2atm'] = df_sol['CO2_PAL'] * pco2_i * 1e6
        df_sol['O2PAL'] = df_sol['O']/o0
        df_sol.to_csv("tracer.csv")
        df_params.to_csv("parameters.csv")

        fig, axes = plt.subplots(figsize = (12,10), nrows = 3, ncols = 3)

        raw_pco2.plot(x='age', y='pco2', ax=axes[0,0], marker = '*', lw=0)
        df_sol.plot(x='Age', y='CO2atm', ax=axes[0,0])

        df_sol.plot(x='Age', y='U_m', ax=axes[0,1])

        raw_u.plot(x='age', y = 'u', ax = axes[0,2], marker='*', lw=1.5)
        df_sol.plot(x='Age', y='d235U', ax=axes[0,2])

        df_sol.plot(x='Age', y='phosphate_m', ax=axes[1,0])
        df_params.plot(x='Age', y ='ANOX', ax=axes[1,1])
        # axes[1,2].remove()
        df_params.plot(x='Age', y ='mocb', ax=axes[1,2])
        df_params.plot(x='Age', y ='locb', ax=axes[1,2])
        # df_params.plot(x='Age', y ='mocb', ax=axes[2,1])
        df_params.plot(x='Age', y ='oxidw', ax=axes[1,2])
        # df_params.plot(x='Age', y ='carbw', ax=axes[2,2])

        raw_d13c.plot(x='age', y='d13c', ax=axes[2,0], marker = '*', lw=1.5)
        df_sol.plot(x='Age', y='d13c', ax=axes[2,0])

        df_sol.plot(x='Age', y='O2PAL', ax=axes[2,1])

        axes[2,2].plot(t, fcinp(t)/1e15)

        axes[2,2].set_ylabel('Carbon emission (Gt/yr)')


        plt.tight_layout()

        pco2_proxy = target['pco2'].values
        pco2_std = target['pco2_std'].values
        d13c_proxy = target['d13c'].values
        d13c_std = target['d13c_std'].values
        u_proxy = target['u'].values
        u_std = target['u_std'].values

        pco2_model = df_sol['CO2atm'].values
        d13c_model = df_sol['d13c'].values
        u_model = df_sol['d235U'].values

        # u_index1 = range(21)
        # u_index2 = 35

        # index = [0,5,11,15,20,37, 44,50,55,60,65,70,73]

        # sum_diff = sum((pco2_proxy-pco2_model)**2/(pco2_std**2)) +   sum((u_proxy[0:-1]-u_model[0:-1])**2/(u_std[0:-1]**2)) + sum((d13c_proxy[0:-1]-d13c_model[0:-1])**2/(d13c_proxy[0:-1]**2))
        # sum_diff += sum(np.log(2*np.pi*pco2_std**2))+sum(np.log(2*np.pi*u_std[0:-1]**2))+sum(np.log(2*np.pi*d13c_std[0:-1]**2))



        # return lp-0.5 * sum_diff
        return 0

# labels = ["ppal_i", "scale_silw1", "scale_silw2","scale_silw3", "scale_locb", 'scale_degassing', 'scale_u_riv', 'cinput1', 'cinput2', 'cinput3', 'cinput4', 'cinput5', 'd13c1', 'd13c2', 'd13c3', 'd13c4', 'd13c5']

params_test = [  1.23984592,   1.897619,   2.5,   1.3559425,   1.07638667,
   0.86314637,   0.79446403,   
#                0,0,0,0,0,
    0,   5.16283163,   5.44446281, 7.85036545,  20,
    -19.71662526, 2,  -6.75820334, -6.00757177,  -6.02946557]
a = np.copy(param_50_percent)

a[10]=20
a[7] =6
print(a)
ode_plot(a)

[  1.24784648   1.44319999   1.99779317   1.36407063   0.98693632
   1.01252511   4.81280269   6.           5.63666363   5.91339294
  20.         -17.80324775  -5.95476183  -7.02727398  -8.55497833
  -6.6900894 ]

0

param_50_percent

array([  1.24784648,   1.44319999,   1.99779317,   1.36407063,
         0.98693632,   1.01252511,   4.81280269,   5.48192546,
         5.63666363,   5.91339294,  30.57730917, -17.80324775,
        -5.95476183,  -7.02727398,  -8.55497833,  -6.6900894 ])